"""
@Description :   读取网络文件，社区检测
@Author      :   Li Junjie 
@Time        :   2025/03/19 20:40:08
"""
import networkx as nx
import community as community_louvain
import matplotlib.pyplot as plt
from pyvis.network import Network
import matplotlib.cm as cm
import matplotlib.colors as colors
import json

# 假设你的网络图为 G（已创建）
G = nx.read_gexf("data/network_graph.gexf")
# 使用 Louvain 算法进行社区检测
partition = community_louvain.best_partition(G.to_undirected(), random_state=42)

# 将社区标签赋给每个节点
nx.set_node_attributes(G, partition, 'community')

# 查看社区数量
num_communities = len(set(partition.values()))
print(f'检测到的社区数量：{num_communities}')

# 为每个社区指定颜色
cmap = cm.get_cmap('tab20', num_communities)
community_colors = {community: colors.to_hex(cmap(community)) for community in set(partition.values())}

# 使用 Pyvis 可视化网络（带社区颜色）
net = Network(height="800px", width="100%", bgcolor="#ffffff", font_color="black", cdn_resources="remote", directed=True)

for node, attrs in G.nodes(data=True):
    community = attrs['community']
    net.add_node(
        node,
        label=str(node),
        color=community_colors[community],
        title=f"节点: {node}<br>社区: {community}<br>海拔: {attrs.get('elevation')}"
    )

# 添加边
for source, target in G.edges():
    net.add_edge(source, target)

# 提供交互调节功能
net.show_buttons(filter_=['physics'])

# 保存并显示网络
net.save_graph("html/network_communities.html")
# net.show("html/network_communities.html", notebook=False)

# 构建节点信息字典（编号、社区、颜色）
node_community_info = {}
for node, community_id in partition.items():
    node_community_info[node] = {
        'community': community_id,
        'color': community_colors[community_id]
    }

# 保存为JSON文件
output_path = 'data/node_community_colors.json'
with open(output_path, 'w', encoding='utf-8') as f:
    json.dump(node_community_info, f, ensure_ascii=False, indent=4)

print(f"节点社区信息已保存到: {output_path}")
